import os
from openai import OpenAI
import json
from typing import List, Dict, Optional

class KimiClient:
    def __init__(self, api_key: Optional[str] = None, base_url: str = "https://api.moonshot.cn/v1"):
        """
        初始化Kimi API客户端
        
        Args:
            api_key: Kimi API密钥，如果为None则从环境变量读取
            base_url: API基础URL
        """
        self.api_key = api_key or os.getenv("KIMI_API_KEY")
        if not self.api_key:
            raise ValueError("API密钥未提供且未设置KIMI_API_KEY环境变量")
        
        self.client = OpenAI(
            api_key=self.api_key,
            base_url=base_url
        )
        self.model = "kimi-k2-0905-preview"
        self.conversation_history: List[Dict] = []
    
    def add_system_message(self, content: str = "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。"):
        """添加系统消息到对话历史"""
        self.conversation_history.append({
            "role": "system",
            "content": content
        })
    
    def chat(self, message: str, temperature: float = 0.6) -> str:
        """
        发送消息给Kimi并获取回复
        
        Args:
            message: 用户消息
            temperature: 生成温度
            
        Returns:
            Kimi的回复内容
        """
        # 添加用户消息到历史
        self.conversation_history.append({
            "role": "user",
            "content": message
        })
        
        try:
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=self.conversation_history,
                temperature=temperature
            )
            
            response = completion.choices[0].message.content
            
            # 添加助手回复到历史
            self.conversation_history.append({
                "role": "assistant",
                "content": response
            })
            
            return response
            
        except Exception as e:
            # 移除最后一条用户消息（因为发送失败）
            self.conversation_history.pop()
            raise Exception(f"API调用失败: {str(e)}")
    
    def clear_history(self):
        """清空对话历史"""
        self.conversation_history = []
    
    def get_history(self) -> List[Dict]:
        """获取对话历史"""
        return self.conversation_history.copy()
    
    def set_history(self, history: List[Dict]):
        """设置对话历史"""
        self.conversation_history = history.copy()
